{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SageMaker V3 Model Optimization Example\n",
    "\n",
    "This notebook demonstrates how to use SageMaker V3 ModelBuilder to optimize a JumpStart model for improved inference performance."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prerequisites\n",
    "Note: Ensure you have sagemaker and ipywidgets installed in your environment. The ipywidgets package is required to monitor endpoint deployment progress in Jupyter notebooks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import required libraries\n",
    "import json\n",
    "import uuid\n",
    "import time\n",
    "import boto3\n",
    "\n",
    "from sagemaker.serve.model_builder import ModelBuilder\n",
    "from sagemaker.serve.builder.schema_builder import SchemaBuilder\n",
    "from sagemaker.core.resources import EndpointConfig\n",
    "from sagemaker.core.helper.session_helper import Session"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 1: Configure Model and Session\n",
    "\n",
    "We'll optimize a Llama 3 model from JumpStart using AWQ quantization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Configuration\n",
    "MODEL_ID = \"meta-textgeneration-llama-3-8b-instruct\"\n",
    "MODEL_NAME_PREFIX = \"jumpstart-optimize-example\"\n",
    "ENDPOINT_NAME_PREFIX = \"jumpstart-optimize-example-endpoint\"\n",
    "AWS_ACCOUNT_ID = \"593793038179\"\n",
    "AWS_REGION = \"us-east-2\"\n",
    "\n",
    "# Generate unique identifiers\n",
    "unique_id = str(uuid.uuid4())[:8]\n",
    "model_name = f\"{MODEL_NAME_PREFIX}-{unique_id}\"\n",
    "endpoint_name = f\"{ENDPOINT_NAME_PREFIX}-{unique_id}\"\n",
    "job_name = f\"js-optimize-{int(time.time())}\"\n",
    "\n",
    "print(f\"Model name: {model_name}\")\n",
    "print(f\"Endpoint name: {endpoint_name}\")\n",
    "print(f\"Optimization job name: {job_name}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 2: Create Schema Builder\n",
    "\n",
    "Define the input/output schema for the text generation model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create schema builder for text generation\n",
    "sample_input = {\"inputs\": \"What are falcons?\", \"parameters\": {\"max_new_tokens\": 32}}\n",
    "sample_output = [{\"generated_text\": \"Falcons are small to medium-sized birds of prey.\"}]\n",
    "\n",
    "schema_builder = SchemaBuilder(sample_input, sample_output)\n",
    "print(\"Schema builder created successfully!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 3: Initialize SageMaker Session\n",
    "\n",
    "Create a SageMaker session with the specified AWS region."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create SageMaker session\n",
    "boto_session = boto3.Session(region_name=AWS_REGION)\n",
    "sagemaker_session = Session(boto_session=boto_session)\n",
    "print(f\"SageMaker session created for region: {AWS_REGION}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 4: Create ModelBuilder\n",
    "\n",
    "Initialize the ModelBuilder with the JumpStart model ID and schema."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize ModelBuilder\n",
    "model_builder = ModelBuilder(\n",
    "    model=MODEL_ID,\n",
    "    schema_builder=schema_builder,\n",
    "    sagemaker_session=sagemaker_session,\n",
    ")\n",
    "print(\"ModelBuilder created successfully!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 5: Optimize the Model\n",
    "\n",
    "Optimize the model using AWQ quantization for improved inference performance. This step may take up to 30 minutes to complete!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Optimize the model with AWQ quantization\n",
    "print(\"Optimizing JumpStart model...\")\n",
    "optimized_model = model_builder.optimize(\n",
    "    instance_type=\"ml.g5.2xlarge\",\n",
    "    image_uri=\"763104351884.dkr.ecr.us-east-2.amazonaws.com/djl-inference:0.31.0-lmi13.0.0-cu124\",\n",
    "    output_path=f\"s3://sagemaker-us-east-2-593793038179/optimize-output/jumpstart-{unique_id}/\",\n",
    "    quantization_config={\"OverrideEnvironment\": {\"OPTION_QUANTIZE\": \"awq\"}},\n",
    "    accept_eula=True,\n",
    "    job_name=job_name,\n",
    "    model_name=model_name,\n",
    ")\n",
    "print(f\"Model Successfully Optimized: {optimized_model.model_name}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 6: Deploy the Optimized Model\n",
    "\n",
    "Deploy the optimized model to a SageMaker endpoint for real-time inference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Deploy the optimized model to an endpoint\n",
    "print(\"Deploying optimized model to endpoint...\")\n",
    "core_endpoint = model_builder.deploy(\n",
    "    endpoint_name=endpoint_name,\n",
    "    initial_instance_count=1\n",
    ")\n",
    "print(f\"Endpoint Successfully Created: {core_endpoint.endpoint_name}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 7: Test the Optimized Endpoint\n",
    "\n",
    "Send a test request to verify the optimized model is working correctly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test optimized model invocation\n",
    "test_data = {\n",
    "    \"inputs\": \"What are the benefits of machine learning?\",\n",
    "    \"parameters\": {\"max_new_tokens\": 50}\n",
    "}\n",
    "\n",
    "result = core_endpoint.invoke(\n",
    "    body=json.dumps(test_data),\n",
    "    content_type=\"application/json\"\n",
    ")\n",
    "\n",
    "response_body = result.body.read().decode('utf-8')\n",
    "prediction = json.loads(response_body)\n",
    "print(f\"Result of invoking optimized endpoint: {prediction}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 8: Clean Up Resources\n",
    "\n",
    "Clean up the created resources to avoid ongoing charges."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Clean up resources\n",
    "core_endpoint_config = EndpointConfig.get(endpoint_config_name=core_endpoint.endpoint_name)\n",
    "\n",
    "# Delete in the correct order\n",
    "optimized_model.delete()\n",
    "core_endpoint.delete()\n",
    "core_endpoint_config.delete()\n",
    "\n",
    "print(\"Optimized model and endpoint successfully deleted!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "This notebook demonstrated:\n",
    "1. Creating a ModelBuilder with a JumpStart model\n",
    "2. Optimizing the model using AWQ quantization\n",
    "3. Deploying the optimized model to a SageMaker endpoint\n",
    "4. Making inference requests to the optimized endpoint\n",
    "5. Cleaning up resources\n",
    "\n",
    "The V3 ModelBuilder's optimize() method makes it easy to improve model performance with quantization and other optimization techniques!"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv-test",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
